using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Globalization;
using System.IO;
using System.Xml;

using Atomic.Libraries;
using Atomic.Structures;
using Atomic.Vasp.Options;

namespace Atomic.Vasp
{
	/// <summary>
	/// Executes the VASP density functional theory implementation.
	/// </summary>
	[Serializable]
	public class VaspProblem
	{
		public VaspProblem()
		{
			Options = new VaspOptionCollection();
			Convergence = new VaspConvergence();
			Parallelization = new VaspParallelization();
			Output = new VaspOutput();
			Potentials = new VaspPotentials();

			// Prefer reproducible output. Unset to restore VASP default behavior.
			RandomSeed = 1;
		}

		/*public void Prepare(DirectoryInfo directory)
		{
			// Just write VASP files to the specified directory (or WorkingDirectory if not specified).
			throw new NotImplementedException();
		}*/

		/*public void Estimate()
		{
			// Run VASP makeparam to estimate memory usage.
			throw new NotImplementedException();
		}*/

		public VaspResult Run()
		{
			if (WorkingDirectory != null)
			{
				if (!WorkingDirectory.Exists)
				{
					WorkingDirectory.Create();
				}
				else
				{
					// Don't allow to overwrite files from a previous run unless specified.
					foreach (string f in new string[] { "CHG", "CHGCAR", "CONTCAR", "DOSCAR", "DYNMAT", "EIGENVAL", "IBZKPT", "INCAR", "KPOINTS", "OSZICAR", "OUTCAR", "PCDAT", "POSCAR", "POTCAR", "WAVECAR", "XDATCAR", "vasprun.xml" })
					{
						// Also look for compressed files.
						foreach (string ext in new string[] { "", ".bz2" } )
						{
							FileInfo file = WorkingDirectory.AddFile(f).AddFileExtension(ext);
							if (file.Exists)
							{
								if (Output.OverwriteMode == VaspOutputOverwriteMode.None)
								{
									throw new Exception("VASP output folder not empty.");
								}

								if (Output.OverwriteMode == VaspOutputOverwriteMode.Overwrite)
								{
									file.Delete();
								}
							}
						}
					}
				}

				return RunDirectory(WorkingDirectory);
			}
			else
			{
				// No directory specified. Run VASP in a temporary directory and remove all output files when VASP finishes (and output is parsed).
				using (TempDirectory temp = new TempDirectory())
				{
					return RunDirectory(temp.Directory);
				}
			}
		}
		
		private VaspResult RunDirectory(DirectoryInfo directory)
		{
			// Additional options are added to an internal copy of the options list.
			VaspOptionCollection options = new VaspOptionCollection(Options);
			Prepare(directory, options);
			WriteIncar(directory, options);

			// Execute VASP here.
			RunVasp(directory);

			// Parse output files.
			VaspResult result = Finalize(directory);

			return result;
		}

		protected virtual void Prepare(DirectoryInfo directory, VaspOptionCollection options)
		{
			// Overview of VASP files:
			// http://www.vasp.at/vasp-workshop/slides/handsonII.pdf

			WriteStructure(directory, options);

			// Not really needed? Prefer a minimal INCAR.
			//options.Add("SYSTEM", Structure.Formula);

			if (Convergence.Precision.HasValue)
			{
				options.Add("PREC", Convergence.Precision.ToString().ToUpper());
			}

			if (Convergence.PlaneWaveCutOff.HasValue)
			{
				options.Add("ENCUT", Convergence.PlaneWaveCutOff.Value);
			}

			if (Convergence.EnergyBreakCondition.HasValue)
			{
				options.Add("EDIFF", Convergence.EnergyBreakCondition.Value);
			}

			/*if (Convergence.MaxIonicSteps.HasValue)
			{
				options.Add("NSW", Convergence.MaxIonicSteps.Value);
			}*/

			if (Convergence.MaxElectronicSteps.HasValue)
			{
				options.Add("NELM", Convergence.MaxElectronicSteps.Value);
			}

			if (Parallelization.Bands.HasValue)
			{
				options.Add("NPAR", Parallelization.Bands.Value);
			}

			if (Parallelization.CoresPerOrbital.HasValue)
			{
				options.Add("NCORE", Parallelization.CoresPerOrbital.Value);
			}

			if (Output.ChargeDensity.HasValue)
			{
				options.Add("LCHARG", Output.ChargeDensity.Value);
			}

			if (Output.ChargeDensityIgnore)
			{
				// Ignores the CHGCAR when parsing output files. This is indicated by creating a CHGCAR.ignore file.
				File.WriteAllText(directory.AddFile("CHGCAR.ignore").FullName, "");
			}

			if (Output.WaveFunctions.HasValue)
			{
				options.Add("LWAVE", Output.WaveFunctions.Value);
			}

			if (RandomSeed.HasValue)
			{
				options.Add("RANDOM_SEED", RandomSeed.Value);
			}
		}

		protected virtual VaspResult Finalize(DirectoryInfo directory)
		{
			return new VaspResult(directory);
		}

		private void WriteIncar(DirectoryInfo directory, VaspOptionCollection options)
		{
			VaspIncar incar = new VaspIncar();
			foreach (VaspOption option in options)
			{
				option.Write(incar, directory);
			}
			incar.Write(directory);
		}

		private void WriteStructure(DirectoryInfo directory, VaspOptionCollection options)
		{
			foreach (Site site in Structure.Sites)
			{
				if (site.Atom is VacancyAtom || site.Atom is PlaceholderAtom)
				{
					throw new ArgumentException("Vacancies and placeholders aren't supported by VASP.");
				}
			}

			// Write POSCAR file.
			IEnumerable<Site> sortedSites;
			VaspPoscar.Write(Structure, new FileInfo(Path.Combine(directory.FullName, "POSCAR")), out sortedSites);

			// Write POTCAR file.
			WritePotcar(directory, sortedSites);

			// Determine magnetic moments. FIXME sorted!!
			bool spinPolarized = SpinPolarized.HasValue && SpinPolarized.Value;
			List<double> magneticMoments = new List<double>();
			foreach (Site site in sortedSites)
			{
				SpinPolarizedAtom pa = site.Atom as SpinPolarizedAtom;
				if (pa != null)
				{
					if (SpinPolarized.HasValue && !SpinPolarized.Value)
					{
						// Explicitly disabled. Fail.
						throw new ArgumentException("Spin-polarized calculation is disabled but required by the structure.");
					}

					spinPolarized = true;
					magneticMoments.Add(pa.MagneticMoment);
				}
				else
				{
					magneticMoments.Add(0.0);
				}
			}

			if (spinPolarized)
			{
				List<string> entries = new List<string>();

				while (magneticMoments.Count > 0)
				{
					double x = magneticMoments[0];

					// Determine number of consecutive identical moments.
					int n = 1;
					while (n < magneticMoments.Count && magneticMoments[n] == x)
					{
						n++;
					}

					if (n == 1)
					{
						entries.Add(x.ToString("r", CultureInfo.InvariantCulture));
					}
					else
					{
						entries.Add(string.Format(CultureInfo.InvariantCulture, "{0}*{1:r}", n, x));
					}

					magneticMoments.RemoveRange(0, n);
				}

				options.Add("ISPIN", "2");
				//options.Add("MAGMOM", string.Join(" ", magneticMoments.ConvertAll<string>(moment => moment.ToString("r", CultureInfo.InvariantCulture)).ToArray()));
				options.Add("MAGMOM", string.Join(" ", entries));
			}
			else if (SpinPolarized.HasValue)
			{
				options.Add("ISPIN", SpinPolarized.Value ? "2" : "1");
			}
		}

		private void WritePotcar(DirectoryInfo directory, IEnumerable<Site> sortedSites)
		{
			HashSet<Atom> atoms = new HashSet<Atom>();

			using (StreamWriter writer = new StreamWriter(Path.Combine(directory.FullName, "POTCAR")))
			{
				foreach (Site site in sortedSites)
				{
					// Write the potential of each atom only once.
					Atom atom = site.Atom;
					if (!atoms.Contains(atom))
					{
						writer.Write(Potentials.Read(atom));
						atoms.Add(atom);
					}
				}
				writer.Close();
			}
		}

		private void RunVasp(DirectoryInfo directory)
		{
			string vaspExecutable;
			if (VaspExecutable != null)
			{
				vaspExecutable = VaspExecutable.FullName;
			}
			else
			{
				// Try obtaining VASP executable file from environment variable. Maintaining compatibility with ASE.
				vaspExecutable = Environment.GetEnvironmentVariable("VASP_SCRIPT") ?? "vasp";
			}

			Process p = new Process();
			p.StartInfo.FileName = vaspExecutable;
			p.StartInfo.WorkingDirectory = directory.FullName;
			p.StartInfo.UseShellExecute = false;

			if (!Output.PrintConsole)
			{
				// Redirect but ignore.
				p.StartInfo.RedirectStandardOutput = true;
				p.StartInfo.RedirectStandardError = true;

				// Empty buffers asynchronously.
				//p.OutputDataReceived += (sender, e) => {};
				//p.ErrorDataReceived += (sender, e) => {};
			}

			p.Start();

			if (!Output.PrintConsole)
			{
				// Empty buffers asynchronously.
				// http://stackoverflow.com/a/7608823
				p.BeginOutputReadLine();
				p.BeginErrorReadLine();
			}

			p.WaitForExit();

			if (Output.Compression)
			{
				// Do this even if VASP failed for some reason.
				CompressOutput(directory);
			}

			if (p.ExitCode == 41)
			{
				// Usually reported as:
				// forrtl: severe (41): insufficient virtual memory
				throw new VaspOutOfMemoryException();
			}

			if (p.ExitCode != 0)
			{
				throw new VaspException(p.ExitCode);
			}
		}

		private void CompressOutput(DirectoryInfo directory)
		{
			foreach (string f in new string[] { "CHG", "CHGCAR", "DYNMAT", "EIGENVAL", "OUTCAR", "POTCAR", "WAVECAR", "XDATCAR", "vasprun.xml" })
			{
				FileInfo file = directory.AddFile(f);
				if (file.Exists && file.Length > 0)
				{
					Process p = new Process();
					try
					{
						p.StartInfo.FileName = "bzip2";
						p.StartInfo.Arguments = f;
						p.StartInfo.WorkingDirectory = directory.FullName;
						p.StartInfo.UseShellExecute = false;
						p.Start();
						p.WaitForExit();
					}
					catch
					{
						// Don't fail on being unable to compress output files. Silently ignoring is fine here.
					}
				}
			}
		}

		public Structure Structure
		{
			get;
			set;
		}

		public VaspOptionCollection Options
		{
			get;
			private set;
		}

		public VaspConvergence Convergence
		{
			get;
			private set;
		}

		public VaspParallelization Parallelization
		{
			get;
			private set;
		}

		public VaspOutput Output
		{
			get;
			private set;
		}

		public VaspPotentials Potentials
		{
			get;
			set;
		}

		/// <summary>
		/// Do spin-polarized calculations.
		/// http://cms.mpi.univie.ac.at/vasp/guide/node99.html
		/// </summary>
		public bool? SpinPolarized
		{
			get;
			set;
		}

		/// <summary>
		/// Random seed for initial wavefunction arrays.
		/// http://cms.mpi.univie.ac.at/vasp/vasp/RANDOM_SEED.html
		/// </summary>
		public int? RandomSeed
		{
			get;
			set;
		}

		/// <summary>
		/// The file that's being executed to solve the problem. If not defined, the environment variable VASP_SCRIPT is used.
		/// If that's not defined either, vasp is searched for by the shell.
		/// </summary>
		public FileInfo VaspExecutable
		{
			get;
			set;
		}

		/// <summary>
		/// Working directory for VASP files. A temporary directory is used if not specified.
		/// </summary>
		public DirectoryInfo WorkingDirectory
		{
			get;
			set;
		}
	}
}
